package com.tuniu.agents.impl.base;

import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.annotation.JsonPropertyDescription;
import com.tuniu.agents.advisor.AgentChatMemoryAdvisor;
import com.tuniu.agents.agent.*;
import com.tuniu.agents.common.AgentsConstants;
import com.tuniu.agents.common.CustomerBeanOutputConverter;
import com.tuniu.agents.common.i18n.I18nConstants;
import com.tuniu.agents.message.MessageSender;
import com.tuniu.agents.util.I18nUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.util.Assert;
import reactor.core.publisher.Flux;
import reactor.core.publisher.FluxSink;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;
import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY;

@Slf4j
public class Coordinator extends OptionsAgent {

    public static final String AGENT_ID = "Coordinator";

    public static final String PROMPT_ID = "coordinator";

    private final Map<String, String> matchAgents;

    private final AgentManager agentManager;

    @Autowired
    private MessageSender messageSender;

    /**
     * {{user_message_current}} It is assembled in the Selection group before coming over.
     * Analyst：
     * The user previously mentioned these requirements:
     * 1: XXX
     * 2: XXX
     * The original user demand is:
     * XXX
     * The current recommended result is:
     * XXX (in Markdown format)
     */
    public static final String MESSAGE_TEMPLATE = """
            {{user_message_his}}
            The original user demand is:
            {{user_message_current}}
            """;

    public Coordinator(AgentManagerGroup agentManagerGroup, Map<String, String> matchAgents) {
        this(agentManagerGroup, AGENT_ID, PROMPT_ID, matchAgents);
    }

    public Coordinator(AgentManagerGroup agentManagerGroup, String agentId, String promptTemplateId, Map<String, String> matchAgents) {
        super(agentManagerGroup.getModelBuilder(),
                agentManagerGroup.getChatMemory(),
                agentManagerGroup.getAdvisors().stream().map(advisor -> advisor instanceof AgentChatMemoryAdvisor ? new AgentChatMemoryAdvisor(agentManagerGroup.getChatMemory(), false, false) : advisor).collect(Collectors.toList()),
                agentManagerGroup.getPromptManager(),
                agentManagerGroup.getChatEnvironmentMemory(),
                new AgentOptions(agentId, promptTemplateId));
        this.agentOptions.getAdvisorParam().put(AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10);
        this.matchAgents = matchAgents;
        this.agentManager = agentManagerGroup.getAgentManager();
    }

    @Override
    public Flux<String> stream(String conversationId, List<Message> messages, Map<String, Object> context) {
        try {
            int index = 1;
            StringBuffer agentSb = new StringBuffer();
            for(Map.Entry<String, String> entry : matchAgents.entrySet()) {
                agentSb.append(index++).append(".").append(entry.getKey()).append(":").append(entry.getValue()).append("\n");
            }
            CoordinatorOut coordinatorOut =
                    chatClient.prompt()
                            .system(s -> s.params(chatEnvironmentMemory.get(conversationId)))
                            .messages(messages)
                            .advisors(a -> a.param(CHAT_MEMORY_CONVERSATION_ID_KEY, conversationId).
                                    param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, agentOptions.getAdvisorParam().get(CHAT_MEMORY_RETRIEVE_SIZE_KEY)).
                                    param(AgentChatMemoryAdvisor.CHAT_MEMORY_AGENT_ID_KEY, this.agentId)
                                    .param(AgentsConstants.AgentsAdvisorKey.USER_MESSAGE_TEMPLATE, MESSAGE_TEMPLATE)
                                    .param(AgentsConstants.AgentsAdvisorKey.ASSISTANT_MESSAGE_HISTORY_NEED, "false")
                                    .param(AgentsConstants.AgentsAdvisorKey.USER_MESSAGE_HISTORY_NEED, "false")
                                    .param("agents_list", agentSb.toString()))
                            .advisors()
                            .call()
                            .entity(new CustomerBeanOutputConverter<>(CoordinatorOut.class));
            final Collection<String> agents = CollectionUtils.intersection(coordinatorOut.agents, matchAgents.keySet());
            Assert.notEmpty(agents, "Model intent recognition result is empty: " + conversationId + messages);
            if(agents.size() == 1) {
                // Only one agent, directly call the specified Agent and return
                String agent = agents.stream().findFirst().get();
                StringBuilder builder = new StringBuilder().append("\nCurrent user's decomposed requirement is:\n").append(coordinatorOut.agentMessages().get(agent));

                return agentManager.getAgent(agent).stream(conversationId, builder.toString());
            } else {
                List<Agent> agentList = new ArrayList<>();
                Collection<String> newAgents = matchAgents.keySet().stream().filter(agents::contains).toList();
                for (String mainIntent : newAgents) {
                    Agent currentMainIntentAgent = agentManager.getAgent(mainIntent);
                    if (currentMainIntentAgent != null) {
                        agentList.add(currentMainIntentAgent);
                    }
                }

                if (agentList.isEmpty()) {
                    return Flux.just("No suitable Agent found");
                }

                // Use Flux.create and recursion to execute sequentially
                return Flux.create(sink -> {
                    // Start execution from the first Agent
                    processAgentListSequentially(agentList, 0, conversationId, messages ,coordinatorOut.agentMessages(), sink, context);
                });
            }
        } catch (Exception e) {
            log.error("Intent recognition failed,{} ", conversationId, e);
        }
        return Flux.just(I18nUtil.getMessage(conversationId,I18nConstants.COMMON_SYSTEM_ERROR_MESSAGE_KEY));
    }

    /**
     * Helper method to process the output of each Agent and handle the next Agent after completion.
     *
     * @param agents List of agents to process
     * @param index Current index in the agent list
     * @param conversationId Conversation ID
     * @param messages List of messages
     * @param stringStringMap Map of agentId to decomposed message
     * @param sink FluxSink for streaming output
     * @param context Context map
     */
    private void processAgentListSequentially(List<Agent> agents, int index,
                                           String conversationId, List<Message> messages,
                                           Map<String, String> stringStringMap, FluxSink<String> sink, Map<String, Object> context) {
        if (index >= agents.size()) {
            sink.complete();
            return;
        }
        String message = messages.get(0).getText();
        Agent currentAgent = agents.get(index);
        StringBuilder builder = new StringBuilder().append(message).append("\nCurrent user's decomposed requirement is:\n").append(stringStringMap.get(currentAgent.getAgentId()));
        Flux<String> agentFlux = currentAgent.stream(conversationId, List.of(new UserMessage(builder.toString())), context);

        agentFlux.subscribe(
                // Output each message from the current Agent
                sink::next,
                // Handle errors
                error -> {
                    log.error("Agent processing error: " + currentAgent.getClass().getSimpleName(), error);
                    // Continue to the next Agent even if an error occurs
                    processAgentListSequentially(agents, index + 1, conversationId, messages,stringStringMap, sink, context);
                },
                // After the current Agent completes, process the next Agent
                () -> {
                    log.info("Agent {} execution completed, continue to the next Agent", currentAgent.getClass().getSimpleName());
                    processAgentListSequentially(agents, index + 1, conversationId, messages,stringStringMap, sink, context);
                }
        );
    }

    /**
     * Build the return message after intent recognition.
     *
     * @param currentMainIntent Collection of recognized main intents
     * @return Message string
     */
    public String buildAfterMessage(Collection<String> currentMainIntent) {
        StringBuilder sb = new StringBuilder();
        sb.append("Recognized that you want to book");
        if (CollectionUtils.isNotEmpty(currentMainIntent)) {
            for (String mainIntent : currentMainIntent) {
                sb.append(" ").append(matchAgents.get(mainIntent));
            }
        }
        return sb.toString();
    }

    /**
     * Dispatcher Bean for Coordinator output.
     *
     * @param agents List of recognized agents
     * @param agentMessages Decomposed message content for each agent
     */
    record CoordinatorOut(
            @JsonProperty(value = "agents")
            @JsonPropertyDescription("agent name list")
            List<String> agents,

            @JsonProperty(value = "agentMessages")
            @JsonPropertyDescription("Decomposed message content for each agent")
            Map<String, String> agentMessages
    ) {

    }


}
